// Copyright (c) 2025 Inochi Amaoto

/* SPDX-License-Identifier: GPL-2.0 */

#ifndef __VIO_ENTITY_H__
#define __VIO_ENTITY_H__

#include "vapp-common.h"

#define vio_vapp_entity_id_cmp(_x, _y) ((_x) < (_y))

static void vio_vapp_entity_free(struct vio_vapp_entity* entity)
{
  if (entity->data)
    kvfree(entity->data);
  kvfree(entity);
}

static inline void vio_vapp_entity_destroy(struct vio_vapp_entity* entity)
{
  vio_vapp_entity_free(entity);
}

static inline void vio_vapp_entity_assign_status(struct vio_vapp_entity* entity, u32 status)
{
  atomic_set(&entity->status, status);
}
static inline void __vio_entity_set_status_mask(struct vio_vapp_entity* entity, u32 status)
{
  atomic_or(status, &entity->status);
}
static inline void __vio_entity_unset_status_mask(struct vio_vapp_entity* entity, u32 status)
{
  atomic_and(~status, &entity->status);
}
static inline u32 vio_vapp_entity_get_status(struct vio_vapp_entity* entity)
{
  return atomic_read(&entity->status);
}

static void* vio_vapp_entity_alloc_data(u64 count)
{
  pr_info("vio-data-alloc: %llu bytes\n", count);
  return vmalloc(count);
}

static struct vio_vapp_entity* vio_vapp_entity_alloc(u64 id, u64 count)
{
  struct vio_vapp_entity* entity = kmalloc(sizeof(struct vio_vapp_entity), GFP_KERNEL);

  if (unlikely(!entity))
    return ERR_PTR(-ENOMEM);

  entity->data = vio_vapp_entity_alloc_data(count);
  if (unlikely(!entity->data)) {
    vio_vapp_entity_free(entity);
    return ERR_PTR(-ENOMEM);
  }

  entity->id = id;
  entity->alloc = count;
  vio_vapp_entity_assign_status(entity, VIO_STATUS_INITAL);
  init_waitqueue_head(&entity->wait);

  return entity;
}

static inline void vio_vapp_entity_set_type(struct vio_vapp_entity* entity, u32 _type)
{
  entity->type = _type;
}

static inline void vio_vapp_entity_set_fid(struct vio_vapp_entity* entity, u64 _id)
{
  entity->fid = _id;
}

static inline void vio_vapp_entity_free_data(struct vio_vapp_entity* entity)
{
  kvfree(entity->data);
  entity->data = NULL;
  entity->alloc = 0;
}


static __always_inline struct vio_vapp_entity* vapp_entity_tree_get(struct vio_vapp_entity_tree* root, u64 id)
{
  unsigned long flags;
  struct rb_node* node = NULL;
  struct vio_vapp_entity* data;

  spin_lock_irqsave(&root->lock, flags);

  node = root->root.rb_node;
  while (node) {
    data = container_of(node, struct vio_vapp_entity, node);

    if (vio_vapp_entity_id_cmp(id, data->id))
      node = node->rb_left;
    else if (vio_vapp_entity_id_cmp(data->id, id))
      node = node->rb_right;
    else {
      spin_unlock_irqrestore(&root->lock, flags);
      return data;
    }
  }

  spin_unlock_irqrestore(&root->lock, flags);

  return NULL;
}

__maybe_unused
static int vapp_entity_tree_insert(struct vio_vapp_entity_tree* root, struct vio_vapp_entity* node)
{
  unsigned long flags;
  struct rb_node **new = NULL, *parent = NULL;
  struct vio_vapp_entity* this;

  spin_lock_irqsave(&root->lock, flags);

  new = &root->root.rb_node;
  while (*new) {
    this = container_of(*new, struct vio_vapp_entity, node);
    parent = *new;

    if (vio_vapp_entity_id_cmp(node->id, this->id))
      new = &((*new)->rb_left);
    else if (vio_vapp_entity_id_cmp(this->id, node->id))
      new = &((*new)->rb_right);
    else {
      spin_unlock_irqrestore(&root->lock, flags);
      return -EEXIST;
    }
  }

  rb_link_node(&node->node, parent, new);
  rb_insert_color(&node->node, &root->root);
  spin_unlock_irqrestore(&root->lock, flags);

  return 0;
}

__maybe_unused
static struct vio_vapp_entity* vapp_entity_tree_get_or_insert(struct vio_vapp_entity_tree* root, struct vio_vapp_entity* node)
{
  unsigned long flags;
  struct rb_node **new = NULL, *parent = NULL;
  struct vio_vapp_entity* this;

  spin_lock_irqsave(&root->lock, flags);

  new = &root->root.rb_node;
  while (*new) {
    this = container_of(*new, struct vio_vapp_entity, node);
    parent = *new;

    if (vio_vapp_entity_id_cmp(node->id, this->id))
      new = &((*new)->rb_left);
    else if (vio_vapp_entity_id_cmp(this->id, node->id))
      new = &((*new)->rb_right);
    else {
      spin_unlock_irqrestore(&root->lock, flags);
      return this;
    }
  }

  rb_link_node(&node->node, parent, new);
  rb_insert_color(&node->node, &root->root);
  spin_unlock_irqrestore(&root->lock, flags);

  return node;
}

static inline struct vio_vapp_entity* vapp_entity_tree_get_or_insert_helper(struct vio_vapp_entity_tree* root, struct vio_vapp_entity_cpu_header* header)
{
  struct vio_vapp_entity* target;
  struct vio_vapp_entity* node = vio_vapp_entity_alloc(header->id, header->alloc);
  if (IS_ERR(node))
    return NULL;

  target = vapp_entity_tree_get_or_insert(root, node);
  if (target != node)
    vio_vapp_entity_free(node);
  else {
    target->type = header->type;
    target->fid = header->fid;
    vio_vapp_entity_init_map(target);
  }

  return target;
}

static inline void __vapp_entity_tree_remove(struct vio_vapp_entity_tree* root, struct vio_vapp_entity* node)
{
  rb_erase(&node->node, &root->root);
}

static inline void vapp_entity_tree_remove(struct vio_vapp_entity_tree* root, struct vio_vapp_entity* node)
{
  unsigned long flags;
  spin_lock_irqsave(&root->lock, flags);
  __vapp_entity_tree_remove(root, node);
  spin_unlock_irqrestore(&root->lock, flags);
}

static inline void vapp_entity_tree_init(struct vio_vapp_entity_tree* tree)
{
  spin_lock_init(&tree->lock);
  tree->root.rb_node = NULL;
}

__maybe_unused
static void vapp_entity_tree_destroy(struct vio_vapp_entity_tree* tree)
{
  struct rb_node* pos;
  unsigned long flags;
  spin_lock_irqsave(&tree->lock, flags);
  for(pos = rb_first(&tree->root); pos; pos = rb_first(&tree->root)) {
    struct vio_vapp_entity* entity = container_of(pos, struct vio_vapp_entity, node);

    __vapp_entity_tree_remove(tree, entity);
    vio_vapp_entity_destroy(entity);
  }

  spin_unlock_irqrestore(&tree->lock, flags);
}

__maybe_unused
static int vio_vapp_entity_split_packet(struct vio_vapp_entity* entity, struct vio_vapp_entity_tree* tree, struct vapp_send_data_packet* packets[])
{
  size_t nr = vio_vapp_packet_nr_pkt(entity->alloc);
  size_t i;

  pr_info("split-packet: packet number: %zu\n", nr);
  pr_info("split-packet: entity->id: %llu\n", entity->id);
  for (i = 0; i < nr; ++i) {
    struct vapp_send_data_packet* packet = vapp_send_data_packet_alloc();
    u64 dataoff = i * VIO_VAPP_PKT_SIZE;
    u64 len = entity->alloc > dataoff + VIO_VAPP_PKT_SIZE ? VIO_VAPP_PKT_SIZE : entity->alloc - dataoff;
    if (!packet) {
      goto failed;
    }

    packet->entity = entity;
    packet->map = vio_vapp_packet_get_map_id(dataoff);
    packet->pkt.buffer = entity->data + dataoff;
    packet->pkt.header = (struct vio_vapp_entity_header) {
      .id = cpu_to_le64(entity->id),
      .fid = cpu_to_le64(entity->fid),
      .map = cpu_to_le64(vio_vapp_packet_get_map_id(dataoff)),
      .type = cpu_to_le32(entity->type),
      .alloc = cpu_to_le32(entity->alloc),
      .offset = cpu_to_le32(dataoff),
      .length = cpu_to_le32(len),
    };
    packet->tree = tree;

    packets[i] = packet;
  }

  return nr;

 failed:
  while (i--)
    vapp_send_data_packet_free(packets[i]);

  return -ENOMEM;
}

#define vio_entity_status_has(_entity, _mask) \
  (vio_vapp_entity_get_status(_entity) & (_mask))

#define vio_entity_status_is(_entity, _mask) \
  (vio_vapp_entity_get_status(_entity) == (_mask))

#define vio_entity_status_is_inital(_entity) \
  vio_entity_status_is(_entity, VIO_STATUS_INITAL)

#define vio_entity_set_status_mapped(_entity) \
  __vio_entity_set_status_mask((_entity), VIO_STATUS_MMAP)
#define vio_entity_unset_status_mapped(_entity) \
  __vio_entity_unset_status_mask((_entity), VIO_STATUS_MMAP)
#define vio_entity_set_status_req(_entity) \
  __vio_entity_set_status_mask((_entity), VIO_STATUS_REQ)
#define vio_entity_unset_status_req(_entity) \
  __vio_entity_unset_status_mask((_entity), VIO_STATUS_REQ)
#define vio_entity_set_status_rsp(_entity) \
  __vio_entity_set_status_mask((_entity), VIO_STATUS_RSP)
#define vio_entity_unset_status_rsp(_entity) \
  __vio_entity_unset_status_mask((_entity), VIO_STATUS_RSP)
#define vio_entity_set_status_req_check(_entity) \
  __vio_entity_set_status_mask((_entity), VIO_STATUS_REQ_CHECK)
#define vio_entity_unset_status_req_check(_entity) \
  __vio_entity_unset_status_mask((_entity), VIO_STATUS_REQ_CHECK)



#endif // ! __VIO_ENTITY_H__
